Skip to content

[ARM] MLAS: portable machine-code SVE elementwise kernels, FEXPA exp - #31145

Draft
mirounga wants to merge 2 commits into
microsoft:mainfrom
mirounga:rs_sve_eltwise
Draft

[ARM] MLAS: portable machine-code SVE elementwise kernels, FEXPA exp#31145
mirounga wants to merge 2 commits into
microsoft:mainfrom
mirounga:rs_sve_eltwise

Conversation

@mirounga

Copy link
Copy Markdown
Contributor

Rework the SVE elementwise kernels (f32: Erf, Logistic, Exp, SumExp, Softmax, LogSoftmax, ReduceMaximum, ReduceMinimumMaximum; fp16: Erf, Gelu, Tanh) so they no longer require an SVE-capable compiler or a Linux-only build, and speed them up while doing so.

Structure: each kernel becomes a relocation-free *Impl function that receives its constants through table pointers (no literal pools, no global data references), fronted by a plain C++ dispatch layer (sve/elementwise_sve_dispatch.cpp) that owns the public entry points and fastpaths. The production kernels ship as portable machine code in the style of Arm's KleidiAI library (aarch64/elementwise_sve_asm.S: raw instruction words via KAI_ASM_INST, which is GAS ".inst" on Linux/macOS and armasm64 "DCD" on Windows via aarch64/kai_asm_macros.h, macro set adopted verbatim from KleidiAI). The SVE intrinsics translation units remain in-tree as the reference implementation and regeneration source, selectable with cmake option onnxruntime_SVE_ELEMENTWISE_ASM=OFF.

The exp path uses the SVE FEXPA instruction (base SVE, no extension required) for the SumExp/exp evaluation, with the input clamped at -88.0f: below -127*ln2 the FEXPA index underflows into the exponent field and would produce NaN.

Two kernels gain platform-dispatch coverage on AArch64/SVE that previously existed only on AMD64/RISCV64: ComputeExpF32Kernel and ReduceMinimumMaximumF32Kernel (MLAS_PLATFORM fields moved to a shared guard; call sites in compute.cpp/quantize.cpp extended with MLAS_USE_SVE).

Windows enablement: runtime detection via
PF_ARM_SVE_INSTRUCTIONS_AVAILABLE (SDK-#ifdef-guarded, Windows 11 24H2+), sources wired through the existing cl /P + armasm64 pipeline, the fp16 kernel dispatch un-gated from !_WIN32 (the NEON fp16 fallbacks remain POSIX-only; on Windows without SVE the routines stay null and callers use their scalar fallbacks). Structurally complete; awaiting validation on Windows-on-ARM SVE hardware.

Measured on Cortex-X925/A725 (SVE VL=128), big-core pinned, 104 median cells, versus the SVE intrinsics kernels currently shipping (real-time medians, lower is better):

ComputeSoftmaxInplace (32 cells) 0.772x (1.30x speedup)
GeluErf fused/unfused (36 cells) 0.795-0.797x
Silu (36 cells) 0.927-0.930x
All 104 cells 0.832x (1.20x speedup)

Correctness: full onnxruntime_mlas_test with the machine-code kernels linked: 33213 passed, 0 failed. The code is fully position-independent and vector-length agnostic (runtime cntw/cnth + whilelo loops).

Description

Motivation and Context

Rework the SVE elementwise kernels (f32: Erf, Logistic, Exp, SumExp,
Softmax, LogSoftmax, ReduceMaximum, ReduceMinimumMaximum; fp16: Erf, Gelu,
Tanh) so they no longer require an SVE-capable compiler or a Linux-only
build, and speed them up while doing so.

Structure: each kernel becomes a relocation-free *Impl function that
receives its constants through table pointers (no literal pools, no global
data references), fronted by a plain C++ dispatch layer
(sve/elementwise_sve_dispatch.cpp) that owns the public entry points and
fastpaths. The production kernels ship as portable machine code in the
style of Arm's KleidiAI library (aarch64/elementwise_sve_asm.S: raw
instruction words via KAI_ASM_INST, which is GAS ".inst" on Linux/macOS
and armasm64 "DCD" on Windows via aarch64/kai_asm_macros.h, macro set
adopted verbatim from KleidiAI). The SVE intrinsics translation units
remain in-tree as the reference implementation and regeneration source,
selectable with cmake option onnxruntime_SVE_ELEMENTWISE_ASM=OFF.

The exp path uses the SVE FEXPA instruction (base SVE, no extension
required) for the SumExp/exp evaluation, with the input clamped at -88.0f:
below -127*ln2 the FEXPA index underflows into the exponent field and
would produce NaN.

Two kernels gain platform-dispatch coverage on AArch64/SVE that previously
existed only on AMD64/RISCV64: ComputeExpF32Kernel and
ReduceMinimumMaximumF32Kernel (MLAS_PLATFORM fields moved to a shared
guard; call sites in compute.cpp/quantize.cpp extended with MLAS_USE_SVE).

Windows enablement: runtime detection via
PF_ARM_SVE_INSTRUCTIONS_AVAILABLE (SDK-#ifdef-guarded, Windows 11 24H2+),
sources wired through the existing cl /P + armasm64 pipeline, the fp16
kernel dispatch un-gated from !_WIN32 (the NEON fp16 fallbacks remain
POSIX-only; on Windows without SVE the routines stay null and callers use
their scalar fallbacks). Structurally complete; awaiting validation on
Windows-on-ARM SVE hardware.

Measured on Cortex-X925/A725 (SVE VL=128), big-core pinned, 104 median
cells, versus the SVE intrinsics kernels currently shipping (real-time
medians, lower is better):

  ComputeSoftmaxInplace (32 cells)      0.772x  (1.30x speedup)
  GeluErf fused/unfused (36 cells)      0.795-0.797x
  Silu (36 cells)                       0.927-0.930x
  All 104 cells                         0.832x  (1.20x speedup)

Correctness: full onnxruntime_mlas_test with the machine-code kernels
linked: 33213 passed, 0 failed. The code is fully position-independent
and vector-length agnostic (runtime cntw/cnth + whilelo loops).
@mirounga
mirounga marked this pull request as draft July 29, 2026 21:53
Follow-up to the SVE elementwise kernels, addressing the findings of
Windows-on-ARM validation (GB10 / Cortex-X925+A725, MSVC + armasm64, SVE
active at runtime: 0 test regressions, 1.20-1.60x on the exp/softmax/erf
family on both CPU clusters, but one real regression).

MlasSveReduceMinimumMaximumF32KernelImpl was measured 2.0-3.9x SLOWER than
the generic NEON kernel (monotonic in N, plateauing — worst on the wide
out-of-order X925), and this path feeds dynamic-quantization scale
selection via MlasFindMinMaxElement. Root cause: it was the one kernel
written as a naive loop — a single min/max accumulator pair processing one
vector per iteration, generating a whilelt predicate and using merging
(_m) operations in the hot loop — while the generic NEON kernel runs four
independent accumulator pairs. (This dispatch slot did not exist on ARM
before the parent commit, so the regression was introduced there, and the
Linux benchmark families never covered it.) Rewritten with four
unpredicated accumulator pairs and a predicated tail, matching the
adjacent ReduceMaximum kernel's structure: now at parity with the generic
kernel at VL=128 (1.00x for N >= 1024, was 3.9x slower at the plateau;
worst case 1.12x at N = 128, sub-nanosecond), with headroom on wider-VL
hardware. Numerics are bit-identical to generic, as before.

Add an SVE f32 Tanh kernel: the same clamped P13/Q6 rational polynomial as
the generic kernel (coefficients shared by symbol from MlasTanhConstants),
in the same table-pointer Impl form as Erf/Logistic. The TanhKernelRoutine
dispatch slot moves to the shared SVE/AMD64/RISCV64 platform block.
Measured 0.748-0.751x of the generic kernel's time across N = 128..256K
(1.33x speedup), with identical accuracy (2.15e-07 max relative error,
same polynomial). Note for benchmark harnesses: BM_Tanh is no longer a
negative control.

EltwiseAdd was also examined and deliberately left on the generic path: it
is a memory-bound streaming add, and SVE at VL=128 issues exactly the same
number of 128-bit memory operations as NEON — no headroom (the validation
report's 1.03x reading on it was binary-layout noise, consistent with its
byte-identical fingerprint).

The frozen machine code is regenerated by sve/gen_sve_asm.py, now checked
in (extended for multiple source TUs, a configurable optimization level and
include paths, and a section-driven parser that tolerates gcc's
mid-function stack-probe labels). Regeneration fidelity: all 12 unchanged
functions are byte-identical to the previously shipped assembly; the file
is now 14 functions / 619 instruction words.

Known numerics note carried from validation: the SVE exp path flushes to
zero for inputs below ~-86.7 (the FEXPA path clamps at -88.0 to avoid
exponent-field underflow), where the generic path still returns denormals.
The absolute error is bounded by ~2.4e-38 and is invisible to softmax sums
and to the unit tests' abs-OR-rel tolerance; documented here explicitly.

Correctness: full onnxruntime_mlas_test with the new kernels linked:
33213 passed, 0 failed; the untouched benchmark families re-measured at
0.99-1.01x (byte-identical regeneration moved nothing else).
@@ -0,0 +1,199 @@
#!/usr/bin/env python3
@@ -0,0 +1,199 @@
#!/usr/bin/env python3


def run(cmd):
result = subprocess.run(cmd, capture_output=True, text=True)
args = ap.parse_args()

symbols = args.symbols.split(",")
src_dir = os.path.dirname(os.path.abspath(args.src[0]))
wanted = [s for s in symbols if s not in functions]
check_relocations(obj, wanted)
for name, body in extract_functions(obj, wanted, allow_missing=True).items():
functions[name] = body
args = ap.parse_args()

symbols = args.symbols.split(",")
src_dir = os.path.dirname(os.path.abspath(args.src[0]))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants